import numpy as np

'''头部姿态检测器'''
class HeadDetector_real():

    def __init__(self):
        self.nod_threshold = 0   #点头瞌睡（一阶差分）阈值
        self.backForth_threshold = 0  #前后倾（一阶差分）阈值
        self.rotate_threshold = 0 #旋转（一阶差分）阈值

        self.prePitch = 0  #前一帧pitch旋转角度
        self.preRoll = 0  #前一帧roll旋转角度
        self.preYaw = 0  #前一帧yaw旋转角度
        self.preDistance = 0  #前一帧的distance

        #时间窗口长度默认为3，可以平滑处理
        self.TW_length = 3
        self.delta_nod_list = []  # 前后倾检测的距离差分的绝对值的时间窗口 (快动作),长度为 FRONT_BEHIND_WL
        self.delta_backforth_list = []  # pitch时间窗口 (慢动作),长度为 NODDING_WINDOW_LENGTH
        self.delta_rotate_y_list = []  # pitch差分的绝对值的时间窗口 (慢动作),长度为 NODDING_WINDOW_LENGTH
        self.delta_rotate_z_list = []  # pitch差分的绝对值的时间窗口 (慢动作),长度为 NODDING_WINDOW_LENGTH

        #self.preDetect = 0  #前一帧的检测结果（3个状态，0:静止，1:点头，2:前后倾，3:正常运动） 未用到，因为不用进行频数统计

    def setHeadPose_params(self,nod_threshold=2, backForth_threshold=1.5, rotate_threshold=5):
        '''
        :param nod_threshold:  pitch差分阈值，用于检测点头
        :param backForth_threshold:  距离差分阈值，用于检测前后倾
        :param rotate_threshold: 正常运动检测阈值，如果pitch，roll，yaw超过该阈值，则标记为运动
        :return:
        '''
        self.nod_threshold = nod_threshold
        self.backForth_threshold = backForth_threshold
        self.rotate_threshold = rotate_threshold

    def get_pitch_roll_yaw(self,detector,frame,det,landmark):
        '''
        获取pitch，roll，yaw三个旋转自由度
        @:param detector: feat.Detector实例
        @:param frame: 当前检测帧
        @:param dets: 当前提取的人脸
        @:param landmark: 当前提取的关键点
        '''
        out = detector.detect_facepose(frame, [[det]], landmark)
        pitch,roll,yaw = tuple(out[0][0][0])  # 头部姿态估计  [pitch, roll, yaw]
        return pitch,roll,yaw

    def get_pitch_roll_yaw_x_y_zShift(self,detector,frame,det,landmark):
        '''
        获取pitch，roll，yaw三个旋转自由度, 以及关于x,z,y轴的平移偏移量
        @:param detector: feat.Detector实例
        @:param frame: 当前检测帧
        @:param dets: 当前提取的人脸
        @:param landmark: 当前提取的关键点
        '''
        out = detector.detect_facepose(frame, [[det]], landmark)
        pitch,roll,yaw = tuple(out[0][0][0])  # 头部姿态估计  [pitch, roll, yaw]
        X_shift, Y_shift, Z_shift = tuple(out[0][0][1])
        return pitch, roll, yaw, X_shift[0], Y_shift[0], Z_shift[0]

    def reset_pre_pitch_yaw_roll_dist(self):
        '''清空前一帧的头部姿态特征值为0'''
        self.prePitch = 0
        self.preYaw = 0
        self.preRoll = 0
        self.preDistance = 0

    def head_action_detect(self,detector,frame,det,landmark):
        '''
        运动状态估计：点头（pitch + pitch差值阈值控制），前后倾（distance阈值控制），正常运动（yaw，roll基本差值阈值控制），静止
        :return: 返回一个元组(0/1/2, 3/4)   0:静止(无旋转)， 1:点头， 2:正常旋转运动  3:前后倾, 4:无前后倾  （两种运动不同）
        '''
        nod_flag = 0
        backforth_flag = 0
        rotate_flag  = 0
        pitch, roll, yaw, x_shift, y_shift, z_shift = self.get_pitch_roll_yaw_x_y_zShift(detector, frame, det, landmark)
        delta_pitch, delta_yaw, delta_roll = abs(pitch - self.prePitch), abs(yaw - self.preYaw), abs(roll - self.preRoll)
        # delta_pitch, delta_yaw, delta_roll = abs(x_shift - self.prePitch), abs(y_shift - self.preYaw), abs(z_shift - self.preRoll)   #用平移向量进行点头和摇头检测效果更好
        '''###############################   头部左右旋转检测   ##############################'''
        if (self.preYaw != 0):  # 前一帧距离不为0时，才计算差分
            if(len(self.delta_rotate_y_list) > self.TW_length):
                del self.delta_rotate_y_list[0]
            self.delta_rotate_y_list.append(delta_yaw)

            if (np.mean(self.delta_rotate_y_list) > self.rotate_threshold):  # 头部前后倾
                rotate_flag = 1
        if (self.preRoll != 0):  # 前一帧距离不为0时，才计算差分
            if (len(self.delta_rotate_z_list) > self.TW_length):
                del self.delta_rotate_z_list[0]
            self.delta_rotate_z_list.append(delta_roll)

            if (np.mean(self.delta_rotate_z_list) > self.rotate_threshold):  # 头部前后倾
                rotate_flag = 1

        '''###############################   点头检测   ##############################'''
        self.prePitch = pitch  # 前一帧pitch旋转角度
        self.preRoll = roll  # 前一帧roll旋转角度
        self.preYaw = yaw  # 前一帧yaw旋转角度
        if (self.prePitch != 0):  # 前一帧距离不为0时，才计算差分
            if (len(self.delta_nod_list) > self.TW_length):
                del self.delta_nod_list[0]
            self.delta_nod_list.append(delta_pitch)

            if (np.mean(self.delta_nod_list) > self.nod_threshold and rotate_flag == 0):  # 头部左右旋转不会伴随着头部点头
                nod_flag = 1

        '''###############################   前后倾检测   ##############################'''
        distance = z_shift
        if (self.preDistance != 0):  # 前一帧距离不为0时，才计算差分
            delta_dist = abs(distance - self.preDistance)
            # print(delta_dist)
            if(delta_dist > self.backForth_threshold):  #头部前后倾
                backforth_flag = 1
        self.preDistance = distance

        return nod_flag,rotate_flag,backforth_flag

    # def head_action_detect1(self,detector,frame,det,landmark):
    #     '''
    #     运动状态估计：点头（pitch + pitch差值阈值控制），前后倾（distance阈值控制），正常运动（yaw，roll基本差值阈值控制），静止
    #     :return: 返回一个元组(0/1/2, 3/4)   0:静止(无旋转)， 1:点头， 2:正常旋转运动  3:前后倾, 4:无前后倾  （两种运动不同）
    #     '''
    #     '''###############################   Distance差分时间窗口   ##############################'''
    #     distance = self.get_face_distance(det)
    #     if (self.preDistance != 0):  # 前一帧距离不为0时，才计算差分
    #         if (len(self.deltaDist_list) >= self.FRONT_BEHIND_WL):
    #             self.deltaDist_list.pop(0)  # 弹出最前一个值
    #         self.deltaDist_list.append(abs(distance - self.preDistance))  # 距离差分绝对值用来分析前后倾动作，不用考虑是前倾还是后倾
    #
    #     deltaDist_mean = 0
    #     if(len(self.deltaDist_list) > 0):
    #         deltaDist_mean = np.mean(self.deltaDist_list)  #距离差分的绝对值的均值
    #     self.preDistance = distance
    #
    #     '''###############################   pitch差分时间窗口   ##############################'''
    #     X, Z, Y = self.get_pitch_roll_yaw(detector, frame, det, landmark)
    #     deltaX, deltaY, deltaZ = abs(X - self.prePitch), abs(Y - self.preYaw), abs(Z - self.preRoll)
    #     self.prePitch = X  # 前一帧pitch旋转角度
    #     self.preRoll = Z  # 前一帧roll旋转角度
    #     self.preYaw = Y  # 前一帧yaw旋转角度
    #
    #     # 加入到deltaPitch_list
    #     if (self.prePitch != 0):  # 前一帧距离不为0时，才计算差分
    #         if (len(self.deltaPitch_list) >= self.NODDING_WINDOW_LENGTH):
    #             self.deltaPitch_list.pop(0)  # 弹出最前一个值
    #             self.pitch_list.pop(0)
    #         self.pitch_list.append(X)
    #         self.deltaPitch_list.append(deltaX)  # 距离差分绝对值用来分析前后倾动作，不用考虑是前倾还是后倾
    #
    #     pitch_mean = 0
    #     deltaPitch_mean = 0
    #     if (len(self.deltaPitch_list) > 0):
    #         deltaPitch_mean = np.mean(self.deltaPitch_list)  # 距离差分的绝对值的均值
    #         pitch_mean = np.mean(self.pitch_list)
    #
    #     '''###############################   头部姿态行为判断（旋转判断 + 平移判断）   ##############################'''
    #     front_behind_flag = (deltaDist_mean > self.deltaDist_threshold and deltaY < self.deltaY_threshold)   # 如果距离差分大于阈值且yaw差分小于deltaY_threshold，则判定为前后倾
    #     # print(f"front_behind_flag = {front_behind_flag}, deltaDist_mean = {deltaDist_mean}, deltaY = {deltaY}")
    #     #1、点头判断 + 有无前后倾判断
    #     if(deltaPitch_mean > self.deltaX_threshold and pitch_mean < 0):
    #         if (front_behind_flag):
    #             return 1, 3
    #         else:
    #             return 1, 4  # 点头运动
    #     #2、旋转运动判断 + 有无前后倾判断
    #     # 正常旋转运动
    #     elif (deltaY > self.action_threshold or deltaZ > self.action_threshold or (deltaX > self.action_threshold and X > 0)):  # 如果yaw，roll差分大于阈值，则判定为正常旋转运动
    #         if (front_behind_flag):
    #             return 2, 3
    #         else:
    #             return 2, 4
    #     else:  #无旋转运动
    #         if (front_behind_flag):
    #             return 0,3
    #         else:
    #             return 0,4



